#include <mm/slab.h>
#include <lib/math.h>

#include <string.h>

#define LOG_TAG "valloc"
#include <log.h>

static char kmem_cache_names[][CACHE_NAME_LEN] = {
    "p_lot_8",    "p_lot_16",   "p_lot_32",   "p_lot_64",
    "p_lot_128",  "p_lot_256",  "p_lot_512",  "p_lot_1k",
    "p_lot_2k",   "p_lot_4k",   "p_lot_8k"
};

static char kmem_cache_names_dma[][CACHE_NAME_LEN] = {
    "p_lot_dma_128", "p_lot_dma_256", "p_lot_dma_512",
    "p_lot_dma_1k",  "p_lot_dma_2k",  "p_lot_dma_4k"
};

#define SIZE_OF(x) (sizeof(x) / sizeof(x[0]))

static kmem_cache_t* caches[SIZE_OF(kmem_cache_names)];
static kmem_cache_t* kmem_caches_dma[SIZE_OF(kmem_cache_names_dma)];

void valloc_init(void) {
    for (size_t i = 0; i < SIZE_OF(kmem_cache_names); i++) {
        size_t size = 1 << (i + 3);
        caches[i] = new_kmem_cache(kmem_cache_names[i], size, size > 1024 ? 8 : 1, 0);
    }

    for (size_t i = 0; i < SIZE_OF(kmem_cache_names_dma); i++) {
        size_t size = 1 << (i + 7);
        kmem_caches_dma[i] = new_kmem_cache(kmem_cache_names_dma[i], size, size > 1024 ? 4 : 1, ALIGN_DMA);
    }
}

/**
 * if called from dma, bit_offset is 7
 * because the first kmem_cache is 128 bytes which is 2^7
 */
static void* do_alloc(size_t size, kmem_cache_t** caches, size_t n, size_t bit_offset) {
    size_t i = ilog2(size);
    i += (size - (1 << i) != 0); // if size is not power of 2 then increment i
    i -= bit_offset;

    if (i >= n) {
        LOGE("size too large: %u\n", size);
        return NULL;
    }

    return take_obj(caches[i]);
}

void* valloc(size_t size) {
    return do_alloc(size, caches, SIZE_OF(kmem_cache_names), 3);
}

void* vzalloc(size_t size) {
    void* p = do_alloc(size, caches, SIZE_OF(kmem_cache_names), 3);
    memset(p, 0, size);
    return p;
}

void* valloc_dma(size_t size) {
    return do_alloc(size, kmem_caches_dma, SIZE_OF(kmem_cache_names_dma), 7);
}

void* vzalloc_dma(size_t size) {
    void* p = do_alloc(size, kmem_caches_dma, SIZE_OF(kmem_cache_names_dma), 7);
    memset(p, 0, size);
    return p;
}

static void do_free(void* ptr, kmem_cache_t** caches, size_t n) {
    for (size_t i = 0; i < n; i++) {
        if (release_obj(caches[i], ptr)) {
            return;
        }
    }
    LOGE("do_free: ptr not found: %p\n", ptr);
}

void vfree(void* ptr) {
    do_free(ptr, caches, SIZE_OF(kmem_cache_names));
}

void vfree_dma(void* ptr) {
    do_free(ptr, kmem_caches_dma, SIZE_OF(kmem_cache_names_dma));
}